use byteorder::{BigEndian, ByteOrder};
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream as AsyncTcpStream};
use tokio::spawn;
use tokio::sync::Mutex;

#[async_trait::async_trait]
trait AsyncMessageHandler: Send + Sync {
    async fn handle(
        &self,
        data: Vec<u8>,
        writer: Option<Arc<Mutex<tokio::net::tcp::OwnedWriteHalf>>>,
    ) -> std::io::Result<()>;
}

// 为闭包自动实现特征
#[async_trait::async_trait]
impl<F, Fut> AsyncMessageHandler for F
where
    F: Fn(Vec<u8>, Option<Arc<Mutex<tokio::net::tcp::OwnedWriteHalf>>>) -> Fut + Send + Sync,
    Fut: std::future::Future<Output = ()> + Send + 'static,
{
    async fn handle(
        &self,
        data: Vec<u8>,
        writer: Option<Arc<Mutex<tokio::net::tcp::OwnedWriteHalf>>>,
    ) -> std::io::Result<()> {
        (self)(data, writer).await;
        Ok(())
    }
}

// 拆分 MessageIO 为 ReadMessage 和 WriteMessage trait
#[async_trait::async_trait]
pub trait ReadMessage {
    async fn read_message(&mut self) -> std::io::Result<Vec<u8>>;
}

#[async_trait::async_trait]
pub trait WriteMessage {
    async fn send_message(&mut self, message: &[u8]) -> std::io::Result<()>;
}

// 为 AsyncTcpStream 实现 ReadMessage 和 WriteMessage
#[async_trait::async_trait]
impl ReadMessage for AsyncTcpStream {
    async fn read_message(&mut self) -> std::io::Result<Vec<u8>> {
        let mut size_buffer = [0u8; 4];
        self.read_exact(&mut size_buffer).await?;
        let size = BigEndian::read_u32(&size_buffer) as usize;

        let mut buffer = vec![0u8; size];
        self.read_exact(&mut buffer).await?;
        Ok(buffer)
    }
}

#[async_trait::async_trait]
impl WriteMessage for AsyncTcpStream {
    async fn send_message(&mut self, message: &[u8]) -> std::io::Result<()> {
        let len = message.len() as u32;
        let mut header = [0u8; 4];
        BigEndian::write_u32(&mut header, len);

        self.write_all(&header).await?;
        self.write_all(message).await?;
        self.flush().await?;
        Ok(())
    }
}

// 为 split 后的 OwnedReadHalf 和 OwnedWriteHalf 实现对应 trait
#[async_trait::async_trait]
impl ReadMessage for tokio::net::tcp::OwnedReadHalf {
    async fn read_message(&mut self) -> std::io::Result<Vec<u8>> {
        let mut size_buffer = [0u8; 4];
        self.read_exact(&mut size_buffer).await?;
        let size = BigEndian::read_u32(&size_buffer) as usize;

        let mut buffer = vec![0u8; size];
        self.read_exact(&mut buffer).await?;
        Ok(buffer)
    }
}

#[async_trait::async_trait]
impl WriteMessage for tokio::net::tcp::OwnedWriteHalf {
    async fn send_message(&mut self, message: &[u8]) -> std::io::Result<()> {
        let len = message.len() as u32;
        let mut header = [0u8; 4];
        BigEndian::write_u32(&mut header, len);

        self.write_all(&header).await?;
        self.write_all(message).await?;
        self.flush().await?;
        Ok(())
    }
}

pub struct TcpServer {
    message_handler: Arc<dyn AsyncMessageHandler>,
    port: u16,
}

impl TcpServer {
    pub fn new<F, Fut>(port: u16, handler: F) -> Self
    where
        F: Fn(Vec<u8>, Option<Arc<Mutex<tokio::net::tcp::OwnedWriteHalf>>>) -> Fut
            + Send
            + Sync
            + 'static,
        Fut: std::future::Future<Output = ()> + Send + 'static,
    {
        TcpServer {
            message_handler: Arc::new(handler),
            port,
        }
    }

    pub async fn start_listen(&self) -> std::io::Result<()> {
        let listener = TcpListener::bind(format!("0.0.0.0:{}", self.port)).await?;
        let handler = self.message_handler.clone();
        println!("已成功开始监听");
        loop {
            let (stream, addr) = listener.accept().await?;
            println!("New connection: {}", addr);

            let handler_clone = handler.clone();
            spawn(async move {
                if let Err(e) = handle_client(stream, handler_clone).await {
                    eprintln!("Error handling client: {}", e);
                }
            });
        }
    }
}

async fn handle_client(
    mut stream: AsyncTcpStream,
    handler: Arc<dyn AsyncMessageHandler>,
) -> std::io::Result<()> {
    let (mut reader, writer) = stream.into_split();
    let writer = Arc::new(Mutex::new(writer));
    loop {
        match reader.read_message().await {
            Ok(message) => {
                if message.is_empty() {
                    println!("Client disconnected gracefully");
                    handler.handle(Vec::new(), Some(writer.clone())).await?;
                    break;
                }
                handler.handle(message, Some(writer.clone())).await?;
            }
            Err(e) => {
                if e.kind() == std::io::ErrorKind::UnexpectedEof {
                    println!("Client disconnected");
                } else {
                    eprintln!("Failed to read from connection: {}", e);
                }
                // 触发断开信号
                handler.handle(Vec::new(), Some(writer.clone())).await?;
                break;
            }
        }
    }
    Ok(())
}

pub struct TcpClient {
    writer: Arc<tokio::sync::Mutex<tokio::net::tcp::OwnedWriteHalf>>,
    message_handler: Arc<dyn AsyncMessageHandler>,
}

impl TcpClient {
    pub async fn new<F, Fut>(ip: &str, port: u16, handler: F) -> std::io::Result<Self>
    where
        F: Fn(Vec<u8>, Option<Arc<Mutex<tokio::net::tcp::OwnedWriteHalf>>>) -> Fut
            + Send
            + Sync
            + 'static,
        Fut: std::future::Future<Output = ()> + Send + 'static,
    {
        let addr: SocketAddr = format!("{}:{}", ip, port).parse().unwrap();
        println!("addr: {}", addr);
        let stream = AsyncTcpStream::connect(addr).await?;
        let (reader, writer) = stream.into_split();

        let client = TcpClient {
            writer: Arc::new(tokio::sync::Mutex::new(writer)),
            message_handler: Arc::new(handler),
        };

        client.start_receiver(reader).await;
        Ok(client)
    }

    pub async fn send(&self, data: &[u8]) -> std::io::Result<()> {
        let mut writer = self.writer.lock().await;
        writer.send_message(data).await
    }

    async fn start_receiver(&self, mut reader: tokio::net::tcp::OwnedReadHalf) {
        let handler = self.message_handler.clone();
        spawn(async move {
            loop {
                match reader.read_message().await {
                    Ok(message) => {
                        if message.is_empty() {
                            println!("Server disconnected");
                            handler.handle(Vec::new(), None).await;
                            break;
                        }
                        // 客户端不需要通过回调函数实现消息传递
                        if let Err(e) = handler.handle(message, None).await {
                            eprintln!("Error handling message: {}", e);
                        }
                    }
                    Err(e) => {
                        if e.kind() == std::io::ErrorKind::UnexpectedEof {
                            println!("Server disconnected");
                        } else {
                            eprintln!("Error reading message: {}", e);
                        }
                        handler.handle(Vec::new(), None).await;
                        break;
                    }
                }
            }
            // 通知断开连接
        });
    }
}
